from __future__ import annotations
import inspect
from queue import Queue
import time
from weakref import WeakValueDictionary, WeakSet
from pathlib import Path
from typing import Literal, Optional, Union, Set
from logging import Logger, Formatter, StreamHandler, FileHandler
from logging.handlers import RotatingFileHandler, QueueHandler, QueueListener
import logging
import sys
import gc


# %% Get Manager Instance
_logger_manager = logging.getLoggerClass().manager


# %% File Handlers Manager (Singleton)
# 管理文件句柄, 通过QueueHandler实现异步日志处理
class FileHandlerManager:
    _instance = None
    _handlers: WeakValueDictionary  # name -> filehandler
    _handlers_set: Set  # filehandler
    _queues: Queue  # listener queue
    _listeners: QueueListener
    _queue_handler: QueueHandler

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._handlers = WeakValueDictionary()
            cls._handlers_set = set()
            cls._queues = Queue(-1)
            cls._queue_handler = QueueHandler(cls._queues)
            cls._queue_handler.setFormatter(file_formatter)
            cls._listener = cls.start_listener(cls._queues, cls._handlers_set)
        return cls._instance

    def __contains__(self, query: Union[str, FileHandler]) -> bool:
        if isinstance(query, str):
            return query in self._handlers.keys()
        elif isinstance(query, FileHandler):
            return query in self._handlers.values()

    @staticmethod
    def start_listener(
        queue: Queue,
        handlers_set: WeakSet[FileHandler],
    ) -> QueueListener:
        listener = QueueListener(queue, respect_handler_level=True)
        listener.handlers = handlers_set  # 绑定线程监听对象
        listener.start()  # 自动启动监听线程
        return listener

    def getHandler(
        self,
        handler:FileHandler,
    ) -> QueueHandler:
        key = handler.baseFilename
        if key in self._handlers:
            assert isinstance(handler, type(self._handlers[key])), (
                f"exists handler {key} is not {type(handler)}, got {type(self._handlers[key])}"
            )
            print("[FileHandlerManager] Reusing existing handler:", key)
        else:
            self._handlers_set.add(handler)
            self._handlers[key] = handler
        return self._queue_handler

    def removeHandler(
        self,
        handler: FileHandler,
    ) -> None:
        key = handler.baseFilename
        if key in self._handlers:
            self._handlers_set.discard(handler)
            gc.collect()
        else:
            raise ValueError(f"Handler {key} not found in handlers set")


# %% Formatter
class ColoredFormatter(Formatter):
    COLORS = {
        "WARNING": "\033[33m",  # Yellow
        "ERROR": "\033[31m",  # Red
        "CRITICAL": "\033[31m",  # Red
        "DEFAULT": "\033[32m",  # Green
        "INFO": "\033[34m",  # Blue
    }
    RESET = "\033[0m"
    NAME_WIDTH = 0

    def formatTime(self, record, datefmt=None):
        ct = time.localtime(record.created)
        ms = int((record.created - int(record.created)) * 1e3)
        if datefmt:
            s = time.strftime(datefmt, ct)
            s = f"{s}.{ms:03d}"
        else:
            s = time.strftime("%Y-%m-%d %H:%M:%S", ct)
            s = f"{s}.{ms:03d}"
        return s

    def format(self, record):
        color = self.COLORS.get(record.levelname, self.COLORS["DEFAULT"])
        record.message = record.getMessage().strip()
        record.asctime = self.formatTime(record, self.datefmt)
        self.__class__.NAME_WIDTH = max(self.__class__.NAME_WIDTH, len(record.name))
        s = f"{record.asctime} | {color}{record.levelname:<8}{self.RESET} | {record.name:<{self.__class__.NAME_WIDTH}} | {record.message}"
        return s


class FileFormatter(Formatter):
    NAME_WIDTH = 0

    def formatTime(self, record, datefmt=None):
        ct = time.localtime(record.created)
        ms = int((record.created - int(record.created)) * 1e3)
        if datefmt:
            s = time.strftime(datefmt, ct)
            s = f"{s}.{ms:03d}"
        else:
            s = time.strftime("%Y-%m-%d %H:%M:%S", ct)
            s = f"{s}.{ms:03d}"
        return s

    def format(self, record):
        record.message = record.getMessage().strip()
        record.asctime = self.formatTime(record, self.datefmt)
        self.__class__.NAME_WIDTH = max(self.__class__.NAME_WIDTH, len(record.name))
        s = f"{record.asctime} | {record.levelname:<8} | {record.name:<{self.__class__.NAME_WIDTH}} | {record.message}"
        return s


# 提前实例化Formatter，避免多次实例化, 且可以被提前修改
colored_formatter = ColoredFormatter()
file_formatter = FileFormatter()


# %% Async Logger
class AsyncLogger(Logger):
    def __call__(self, string: str = "", *, level: int = logging.DEBUG):
        super().log(level, string)

    def addHandler(self, hdlr):
        self.propagate = False
        if isinstance(hdlr, FileHandler):
            handler_manager = FileHandlerManager()
            queue_hdlr = handler_manager.getHandler(hdlr)
            if queue_hdlr in self.handlers:
                print("[AsyncLogger] Handler already exists:", hdlr.baseFilename)
            super().addHandler(queue_hdlr)
        else:
            super().addHandler(hdlr)
        self.propagate = True

    def removeHandler(self, hdlr):
        if isinstance(hdlr, FileHandler):
            handler_manager = FileHandlerManager()
            handler_manager.removeHandler(hdlr)
        else:
            return super().removeHandler(hdlr)

    def addFileHandler(
        self,
        log_file: Union[str, Path],
        type: Literal["base", "rotating", "time"] = "rotating",
        level=logging.NOTSET,
        maxBytes=10485760,
        backupCount=10,
    ):
        if isinstance(log_file, str):
            log_file = Path(log_file).expanduser().resolve()
        log_file = str(log_file)
        if type == "base":
            file_handler = FileHandler(log_file)
        elif type == "rotating":
            file_handler = RotatingFileHandler(
                log_file, maxBytes=maxBytes, backupCount=backupCount
            )
        elif type == "time":
            raise NotImplementedError("Time-based file handler not implemented yet")
        file_handler.setLevel(level)
        return self.addHandler(file_handler)

    def trace(
        self, msg, *args, exc_info=None, stack_info=False, stacklevel=1, extra=None
    ):
        return super().debug(
            msg,
            *args,
            exc_info=exc_info,
            stack_info=stack_info,
            stacklevel=stacklevel,
            extra=extra,
        )

    def error(
        self, msg, *args, exc_info=None, stack_info=False, stacklevel=1, extra=None
    ):
        frame = inspect.currentframe().f_back
        if frame is not None:
            filename = frame.f_code.co_filename
            lineno = frame.f_lineno
            msg = f"{filename}:{lineno} | {msg}"
        else:
            msg = f"<unknown>:0 | {msg}"
        return super().error(
            msg,
            *args,
            exc_info=exc_info,
            stack_info=stack_info,
            stacklevel=stacklevel,
            extra=extra,
        )

    def stopwatch(this, timeout_ms: Optional[int] = None):
        class StopwatchContext:
            def __enter__(self, *args, **kwargs):
                self.timeout_ms = timeout_ms
                self.start_time = time.monotonic()
                return self

            def __exit__(self, exc_type, exc_value, traceback):
                frame = inspect.currentframe().f_back
                filename = frame.f_code.co_filename
                lineno = frame.f_lineno
                end_time = time.monotonic()
                duration_ms = int((float(end_time) - float(self.start_time)) * 1000)
                if self.timeout_ms is not None and duration_ms > self.timeout_ms:
                    this(
                        f"{filename}:{lineno} | TAKE {duration_ms} MS > TIMEOUT {self.timeout_ms} MS",
                        logging.WARNING,
                    )
                if self.timeout_ms is None:
                    this(f"{filename}:{lineno} | TAKE {duration_ms} MS", logging.DEBUG)
                return False

        return StopwatchContext()


# %% Global Logger Instance
_logger_manager.setLoggerClass(AsyncLogger)  # 更改默认Logger类
ROOT_LOGGER = _logger_manager.getLogger("eztrade")

# %% functions
def get_logger(
    name: str, *, level: Union[int, str] = logging.NOTSET, log_file=None
) -> AsyncLogger:
    if name == "eztrade":
        return ROOT_LOGGER
    if isinstance(level, str):
        level = logging._nameToLevel[level.upper()]
    # Check if handlers already exist to prevent duplicates
    logger: AsyncLogger = _logger_manager.getLogger(name)
    if name.startswith("eztrade."):
        logger.parent = ROOT_LOGGER
        logger.propagate = True
    logger.setLevel(logging.DEBUG)  # Set to lowest level to allow all messages through
    if not logger.handlers:
        handler = StreamHandler(sys.stdout)
        for hdlr in ROOT_LOGGER.handlers if name.startswith("eztrade.") else []:
            if isinstance(hdlr, StreamHandler) and hdlr.stream in [sys.stderr, sys.stdout]:
                # 如果ROOT_LOGGER已经有流处理器, 则复用
                hdlr.setFormatter(colored_formatter)
                break
        else:
            handler.setFormatter(colored_formatter)
            handler.setLevel(level)
            logger.addHandler(handler)
    # File handler for file output
    if log_file:
        logger.addFileHandler(log_file, level=level)
    return logger

def get_loggers(prefix:str = "eztrade.") -> list[AsyncLogger]:
    """获取当前所有满足prefix的logger实例"""
    loggers = [
        v
        for k, v in _logger_manager.loggerDict.items()
        if k.startswith(prefix) and isinstance(v, AsyncLogger)
    ]
    return loggers
